package com.example.tcp.client.store.impl;

import com.example.tcp.client.RetryTcpClient;
import com.example.tcp.client.serde.Serde;
import com.example.tcp.client.store.AbstractDequeStore;
import com.google.common.primitives.Longs;
import lombok.extern.slf4j.Slf4j;
import org.rocksdb.*;

import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.atomic.AtomicReference;

@Slf4j
public class RocksDBFileQueueDequeStore<T> extends AbstractDequeStore<T> {

    private static final String COUNT = "count";

    private Serde<T> serde;
    private byte[] name;
    private RocksDB db;
    private byte[] curKey = null;

    private long count = 0;
    private ColumnFamilyHandle countFamilyHandle;

    private CurrentTimeId ids = new CurrentTimeId();

    public RocksDBFileQueueDequeStore(String dir, String name, Serde<T> serde) throws RocksDBException {
        this.serde = serde;
        this.name = name.getBytes();
        ColumnFamilyOptions columnFamilyOptions = new ColumnFamilyOptions();
        ColumnFamilyDescriptor countDescriptor = new ColumnFamilyDescriptor(COUNT.getBytes(), columnFamilyOptions);

        DBOptions options = new DBOptions().setCreateIfMissing(true)
                .setAllowMmapReads(true)
                .setAllowMmapWrites(true)
                .setCreateMissingColumnFamilies(true);
        final List<ColumnFamilyHandle> handles = new ArrayList<>(2);
        this.db = RocksDB.open(options, dir, Arrays.asList(new ColumnFamilyDescriptor(RocksDB.DEFAULT_COLUMN_FAMILY), countDescriptor), handles);
        countFamilyHandle = handles.get(1);

        byte[] value = db.get(countFamilyHandle, COUNT.getBytes());
        if(value != null) {
            count = bytesToLong(value, 0, value.length);
        }
    }

    @Override
    protected boolean _enqueue(T e) {
        try {
            db.put(getKey(ids.nextId()), serde.serialize(e));
            count++;
            db.put(countFamilyHandle, COUNT.getBytes(), longToBytes(count));
            return true;
        } catch (RocksDBException e1) {
            e1.printStackTrace();
        }
        return false;
    }

    @Override
    protected T _dequeue() {
        T v = null;
        try {
            v = get(true);
            if( v == null) {
                do {
                    notEmpty.await();
                } while (running.get() && (v = get(true)) == null);
            }
            if(v != null) {
                count--;
                db.put(countFamilyHandle, COUNT.getBytes(), longToBytes(count));
            }
        } catch (RocksDBException | InterruptedException e) {
            e.printStackTrace();
        }

        return v;
    }

    @Override
    protected void _close() {
        if(countFamilyHandle != null) {
            countFamilyHandle.close();
        }
        if(db != null) {
            db.close();
        }
    }

    @Override
    public boolean isEmpty() {
        return count == 0;
    }

    @Override
    public Object wrap() {
        return db;
    }

    private T get(boolean remove) throws RocksDBException {
        RocksIterator it = db.newIterator(new ReadOptions().setFillCache(false));
        if(curKey == null) {
            it.seek(getKey(0));
        } else {
            it.seek(getKey(bytesToLong(curKey, name.length, 8) + 1));
        }
        if(it.isValid()) {
            curKey = it.key();
            byte[] value = it.value();
            if(remove) {
                db.delete(curKey);
            }
            return serde.deserialize(value);
        }
        return null;
    }

    private byte[] getKey(long id) {
        byte[] array = new byte[8 + name.length];
        System.arraycopy(name, 0, array, 0, name.length);
        byte[] time = longToBytes(id);
        for(int i = name.length, j = 0; i < 13; ++i, j++) {
            array[i] = time[j];
        }
        return array;
    }

    @Override
    public T peek() {
        try {
            return get(false);
        } catch (RocksDBException e) {
            e.printStackTrace();
        }
        return null;
    }

    @Override
    public long size() {
        return count;
    }

    public static class CurrentTimeId {
        private  AtomicReference<Long> currentTime = new AtomicReference<>(System.currentTimeMillis());

        public  Long nextId() {
            return currentTime.accumulateAndGet(System.currentTimeMillis(),(prev, next) -> next > prev ? next : prev + 1);
        }
    }


    public static byte[] longToBytes(long v) {
        ByteBuffer buffer = ByteBuffer.allocate(Long.BYTES);
        buffer.putLong(v);
        return buffer.array();
    }

    public static long bytesToLong(byte[] v, int offset, int size) {
        ByteBuffer buffer = ByteBuffer.allocate(Long.BYTES);
        buffer.put(v, offset, size);
        buffer.flip();//need flip
        return buffer.getLong();
    }
}
